package com.cc.wechat.websocket.interceptor;

import com.cc.wechat.security.JwtTokenProvider;
import lombok.RequiredArgsConstructor;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;

@Component
@RequiredArgsConstructor
public class WebSocketInterceptor implements HandshakeInterceptor {

    private final JwtTokenProvider jwtTokenProvider;

    /**
     * 在握手之前调用
     * @param request 请求
     * @param response 响应
     * @param wsHandler WebSocket处理器
     * @param attributes 属性
     * @return 是否继续握手
     * @throws Exception 异常
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                 WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        // 获取token
        ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) request;
        String token = servletRequest.getServletRequest().getParameter("token");
        
        if (token != null && jwtTokenProvider.validateToken(token)) {
            // 验证token并获取用户ID
            Long userId = jwtTokenProvider.getUserIdFromToken(token);
            // 将用户ID存储在attributes中，供WebSocketHandler使用
            attributes.put("userId", userId);
            return true;
        }
        return false;
    }

    /**
     * 在握手之后调用
     * @param request 请求
     * @param response 响应
     * @param wsHandler WebSocket处理器
     * @param exception 异常
     */
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                             WebSocketHandler wsHandler, Exception exception) {
        // 握手后的处理，通常为空
    }
} 